#!/usr/bin/env python2.7

import unittest
from dnaseqlib import *
from collections import deque

### Utility classes ###

# Maps integer keys to a set of arbitrary values.
class Multidict:
    # Initializes a new multi-value dictionary, and adds any key-value
    # 2-tuples in the iterable sequence pairs to the data structure.
    def __init__(self, pairs=[]):
        self.storage = {}
        for k,v in pairs:
            self.put(k, v)

    # Associates the value v with the key k.
    def put(self, k, v):
        lst = self.__internal_get(k)
        lst.append(v)

    def __internal_get(self, k):
        if not self.storage.has_key(k):
            lst = []
            self.storage[k] = lst
        else:
            lst = self.storage[k]
        return lst

    # Gets any values that have been associated with the key k; or, if
    # none have been, returns an empty sequence.
    def get(self, k):
        return self.__internal_get(k)
    
    def iteritems(self): 
        return  self.storage.iteritems()

# Given a sequence of nucleotides, return all k-length subsequences
# and their hashes.  (What else do you need to know about each
# subsequence?)
def subsequenceHashes(seq, k):
    sub = deque()
    sub_len = 0

    rh = None
    i = 0

    for c in seq:
        if sub_len == k:
            l = sub.popleft()
            sub.append(c)
            rh.slide(l, c)
            yield (sub, rh.current_hash(), i - k + 1)
        elif sub_len == k - 1:
            sub.append(c)
            rh = RollingHash(sub)
            sub_len += 1
            yield (sub, rh.current_hash(), 0)
        else:
            sub.append(c)
            sub_len += 1
        i += 1

    if sub_len < k:
        yield (sub, 0, -1)

# Similar to subsequenceHashes(), but returns one k-length subsequence
# every m nucleotides.  (This will be useful when you try to use two
# whole data files.)
def intervalSubsequenceHashes(seq, k, m):
    raise Exception("Not implemented!")

# Searches for commonalities between sequences a and b by comparing
# subsequences of length k.  The sequences a and b should be iterators
# that return nucleotides.  The table is built by computing one hash
# every m nucleotides (for m >= k).
def getExactSubmatches(a, b, k, m):
    md = Multidict()
    for sub, h, i in subsequenceHashes(a, k):
        if i % m == 0:
            md.put(h, (''.join(sub), i) ) # (subseq, its offset in seq)

    for sub_b, h, b_idx in subsequenceHashes(b, k):
        sub_a_list = md.get(h)
        for sub_a, a_idx in sub_a_list:
            if sub_a != ''.join(sub_b):
                continue
            yield a_idx, b_idx

if __name__ == '__main__':
    if len(sys.argv) != 4:
        print 'Usage: {0} [file_a.fa] [file_b.fa] [output.png]'.format(sys.argv[0])
        sys.exit(1)

    # The arguments are, in order: 1) Your getExactSubmatches
    # function, 2) the filename to which the image should be written,
    # 3) a tuple giving the width and height of the image, 4) the
    # filename of sequence A, 5) the filename of sequence B, 6) k, the
    # subsequence size, and 7) m, the sampling interval for sequence
    # A.
    compareSequences(getExactSubmatches, sys.argv[3], (500,500), sys.argv[1], sys.argv[2], 8, 100)
